package connection;


import java.net.InetSocketAddress;
import java.security.KeyPair;

import messages.CcAuthNounceMessage;
import messages.CcPKExchangeMessage;
import messages.CcSKExchangeMessage;
import messages.CertChallengeMessage;
import messages.CertSKExchangeMessage;
import messages.SecChannelChallengeMessage;
import messages.UnexpectedMessage;
import messages.UserPassPKExchangeMessage;
import messages.UserPassSKExchangeMessage;

import org.bouncycastle.crypto.CryptoException;

import security.SecureRandomNumber;
import security.crypto.AssimEncryption;
import vericopServer.VericopData;
import client.ClientData;
import connection.exception.ConnectionResetException;

public class TCPSecureClientSocket extends TCPSecureSocket{

	private KeyPair _assimKeyPair;

	/**
	 * New Constructor to make a socket from an address directly.
	 * 
	 * @param iAddr InetAddress item, that contains the target's IP address
	 * @param port The port from the target
	 */
	public TCPSecureClientSocket(InetSocketAddress sockAddr) {
		_tcpSocket = new TCPSocket(sockAddr);
	}

	public TCPSecureClientSocket(InetSocketAddress sockAddr, int bindPort) throws ConnectionResetException{
		_tcpSocket = new TCPSocket(sockAddr, bindPort);

		if(_bypass)
			return;

		_endpointId = "VERICOP";
		_remoteEndpointId = "CERTCOP";

		_establishingSecureChannel = true;
		_secureChannelEstablished = false;

		_remotePublicKey = VericopData.getInstance().getKeyStoreManager().getCertcopPublicKey();

		for(int numRetries = 0; numRetries < _numRetries; numRetries++){
			_secureChannelEstablished = establishCertificateSecureChannel();

			if(_secureChannelEstablished)
				break;
		}

	}

	public TCPSecureClientSocket(InetSocketAddress sockAddr, EstablishmentModes establMode, String username, String password) throws ConnectionResetException{
		_tcpSocket = new TCPSocket(sockAddr);

		if(_bypass)
			return;

		_endpointId = username;
		_remoteEndpointId = "CERTCOP";

		_establishingSecureChannel = true;
		_secureChannelEstablished = false;

		_remotePublicKey = ClientData.getInstance().getKeyStoreManager().getCertcopPublicKey();

		for(int numRetries = 0; numRetries < _numRetries; numRetries++){

			switch(establMode){
			case UsernamePassword: _secureChannelEstablished = establishUserPassSecureChannel(username, password);break;
			case CitizenCard: _secureChannelEstablished = establishCcSecureChannel(username); break;
			}
			
			if(_secureChannelEstablished)
				break;
		}
	}



	private boolean establishUserPassSecureChannel(String username, String password) throws ConnectionResetException{
		_assimKeyPair = AssimEncryption.generateKeyPair();
		_privateKey = _assimKeyPair.getPrivate();

		UserPassPKExchangeMessage exchangePk = new UserPassPKExchangeMessage(username, password, _assimKeyPair.getPublic());
		_channelLastMessage = exchangePk;
		send(exchangePk);

		Object obj = receiveUserPassProtocolMessage(UserPassProtocolPhases.WaitingSessionKey);
		if(obj==null)
			return false;

		_simSessionKey = ((UserPassSKExchangeMessage) obj).getSessionKey(password, _assimKeyPair.getPrivate());

		return channelEstablishmentChallenge();
	}

	private boolean establishCcSecureChannel(String username) throws ConnectionResetException{
		_assimKeyPair = AssimEncryption.generateKeyPair();
		_privateKey = _assimKeyPair.getPrivate();

		CcPKExchangeMessage pkExchange = new CcPKExchangeMessage(username, _assimKeyPair.getPublic());
		_channelLastMessage = pkExchange;
		send(pkExchange);

		CcSKExchangeMessage skMessage = (CcSKExchangeMessage)receiveCcProtocolMessage(CcProtocolPhases.WaitingSessionKey);
		if(skMessage == null)
			return false;

		_simSessionKey = skMessage.getSessionKey(_assimKeyPair, _remotePublicKey);

		if(_simSessionKey == null)
			return false;

		int nounce = SecureRandomNumber.generateNonce();
		CcAuthNounceMessage authNounceReq = new CcAuthNounceMessage(nounce, _simSessionKey, _remotePublicKey);
		_channelLastMessage = authNounceReq;
		send(authNounceReq);

		CcAuthNounceMessage authNounceResp = (CcAuthNounceMessage) receiveCcProtocolMessage(CcProtocolPhases.WaitingAuthNounceAnswer);
		if(authNounceResp == null)
			return false;

		try {
			int receivedNounce = authNounceResp.getNounce(_simSessionKey, _assimKeyPair.getPrivate());

			if(receivedNounce == nounce + 1)
				return channelEstablishmentChallenge();
			else
				return false;

		} catch (CryptoException e) {
			return false;
		}		
	}

	private boolean establishCertificateSecureChannel() throws ConnectionResetException {
		_assimKeyPair = VericopData.getInstance().getKeyStoreManager().getVericopKeyPair();
		_privateKey = _assimKeyPair.getPrivate();

		int nounce = SecureRandomNumber.generateNonce();
		CertChallengeMessage challengeMessage = new CertChallengeMessage(nounce, _remotePublicKey);
		_channelLastMessage = challengeMessage;
		send(challengeMessage);

		CertSKExchangeMessage sessionKeyMessage = (CertSKExchangeMessage) receiveCertificateProtocolMessage(CertificateProtocolPhases.WaitingSessionKey);
		if(sessionKeyMessage == null)
			return false;

		_simSessionKey = sessionKeyMessage.getSessionKey(_assimKeyPair.getPrivate());		
		int serverChallenge = sessionKeyMessage.getClientChallenge();

		int clientChallenge = SecureRandomNumber.generateNonce();
		SecChannelChallengeMessage challengePair = new SecChannelChallengeMessage(clientChallenge, serverChallenge, _simSessionKey);
		_channelLastMessage = challengePair;
		send(challengePair);

		SecChannelChallengeMessage serverAnswer = (SecChannelChallengeMessage)receiveChallengePhaseMessage(ChallengePhases.WaitingServerChallengeAnswer);
		if(serverAnswer == null)
			return false;

		try {
			int serverChallengeAnswer = serverAnswer.getChallenge(_simSessionKey);
			if(serverChallengeAnswer == clientChallenge){
				_establishingSecureChannel = false;
				return true;
			}
		} catch (CryptoException e) {
			return false;
		}
		return false;
	}

	public Object receive() throws ConnectionResetException{
		return super.receive();
	}
	public void send(Object obj) throws ConnectionResetException{
		super.send(obj);
	}

	public boolean isConnected(){
		return super.isConnected();
	}

	public boolean isSecureChannelEstablished(){
		return super.isSecureChannelEstablished();
	}

	public void close(){
		_tcpSocket.close();
	}

	private Object receiveUserPassProtocolMessage(UserPassProtocolPhases phase) throws ConnectionResetException{
		Object obj = null;

		for(int index = 0; index< 6; index++){
			obj = receive();

			if(obj!=null)
				break;
		}

		if(obj == null || obj.getClass().getSimpleName().equals("UnexpectedMessage")){
			return null;
		}else if( obj.getClass().getSimpleName().equals("UserPassSKExchangeMessage")&& phase == UserPassProtocolPhases.WaitingSessionKey){
			return obj;
		}else{
			send(new UnexpectedMessage());
			return null;
		}		
	}

	private Object receiveCcProtocolMessage(CcProtocolPhases phase) throws ConnectionResetException{
		Object obj = null;

		for(int index = 0; index< 6; index++){
			obj = receive();

			if(obj!=null)
				break;
		}

		if(obj == null || obj.getClass().getSimpleName().equals("UnexpectedMessage")){
			return null;
		}else if(obj.getClass().getSimpleName().equals("CcSKExchangeMessage")&& phase == CcProtocolPhases.WaitingSessionKey){
			return obj;
		}else if(obj.getClass().getSimpleName().equals("CcAuthNounceMessage")&& phase == CcProtocolPhases.WaitingAuthNounceAnswer){
			return obj;
		}else{
			send(new UnexpectedMessage());
			return null;
		}		
	}

	private Object receiveCertificateProtocolMessage(CertificateProtocolPhases phase) throws ConnectionResetException{
		Object obj = null;

		for(int index = 0; index< 6; index++){
			obj = receive();

			if(obj!=null)
				break;
		}

		if(obj == null || obj.getClass().getSimpleName().equals("UnexpectedMessage")){
			return null;
		}else if(obj.getClass().getSimpleName().equals("CertSKExchangeMessage")&& phase == CertificateProtocolPhases.WaitingSessionKey){
			return obj;
		}else{
			send(new UnexpectedMessage());
			return null;
		}		
	}

	private Object receiveChallengePhaseMessage(ChallengePhases phase) throws ConnectionResetException{
		Object obj = null;

		for(int index = 0; index< 6; index++){
			obj = receive();

			if(obj!=null)
				break;
		}

		if(obj == null || obj.getClass().getSimpleName().equals("UnexpectedMessage")){
			return null;
		}else if( obj.getClass().getSimpleName().equals("SecChannelChallengeMessage")&& phase == ChallengePhases.WaitingChallengePair){
			return obj;
		}else if( obj.getClass().getSimpleName().equals("SecChannelChallengeMessage")&& phase == ChallengePhases.WaitingServerChallengeAnswer){
			return obj;
		}else{
			send(new UnexpectedMessage());
			return null;
		}		
	}

	private boolean channelEstablishmentChallenge() throws ConnectionResetException{
		Integer clientChallengeValue = SecureRandomNumber.generateNonce();
		SecChannelChallengeMessage clientChallenge = new SecChannelChallengeMessage(clientChallengeValue, _simSessionKey);
		_channelLastMessage = clientChallenge;
		send(clientChallenge);

		SecChannelChallengeMessage serverChallenge = (SecChannelChallengeMessage)receiveChallengePhaseMessage(ChallengePhases.WaitingChallengePair);
		if(serverChallenge == null)
			return false;

		try {
			int serverChallengeValue = serverChallenge.getChallenge(clientChallengeValue, _simSessionKey);
			SecChannelChallengeMessage serverResponse = new SecChannelChallengeMessage(serverChallengeValue, _simSessionKey);
			_channelLastMessage = serverResponse;
			send(serverResponse);

			_establishingSecureChannel = false;
			return true;			
		} catch (CryptoException e) {
			return false;
		}
	}
}
